package org.example.yuaiagent.chatmemory;

import com.alibaba.fastjson.JSON;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.Message;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.stereotype.Component;

import java.io.*;
import java.util.ArrayList;
import java.util.List;

/**
 * 基于文件持久化的对话记忆
 */

@Component
public class FileBasedChatMemory implements ChatMemory {

    @Autowired
    private StringRedisTemplate redisTemplate;

    private final String BASE_DIR;

    // 构造对象时，指定文件保存目录
//    public FileBasedChatMemory(String dir) {
//        this.BASE_DIR = dir;
//        File baseDir = new File(dir);
//        if (!baseDir.exists()) {
//            baseDir.mkdirs();
//        }
//    }

    public FileBasedChatMemory() {
        this.BASE_DIR = System.getProperty("user.dir") + "/tmp/chat-memory";
        File baseDir = new File(this.BASE_DIR);
        if (!baseDir.exists()) {
            baseDir.mkdirs();
        }
    }

    @Override
    public void add(String conversationId, List<Message> messages) {
        List<Message> conversationMessages = getOrCreateConversation(conversationId);
        conversationMessages.addAll(messages);
        saveConversation(conversationId, conversationMessages);
    }

    @Override
    public List<Message> get(String conversationId, int lastN) {
        List<Message> allMessages = getOrCreateConversation(conversationId);
        return allMessages.stream()
                .skip(Math.max(0, allMessages.size() - lastN))
                .toList();
    }

    @Override
    public void clear(String conversationId) {
        File file = getConversationFile(conversationId);
        if (file.exists()) {
            file.delete();
        }
    }

    private List<Message> getOrCreateConversation(String conversationId) {
        File file = getConversationFile(conversationId);
        List<Message> messages = new ArrayList<>();
        if (file.exists()) {
            try (Reader reader = new FileReader(file)) {
                String json = readAll(reader);
                messages = JSON.parseObject(json, ArrayList.class);
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
        return messages;
    }

    private void saveConversation(String conversationId, List<Message> messages) {
        File file = getConversationFile(conversationId);
        try (Writer writer = new FileWriter(file)) {
            String json = JSON.toJSONString(messages);
            redisTemplate.opsForValue().set(conversationId, json);
            writer.write(json);
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    private File getConversationFile(String conversationId) {
        return new File(BASE_DIR, conversationId + ".json"); // 更改为 .json 扩展名
    }

    private String readAll(Reader reader) throws IOException {
        StringBuilder sb = new StringBuilder();
        char[] buffer = new char[1024];
        int read;
        while ((read = reader.read(buffer)) != -1) {
            sb.append(buffer, 0, read);
        }
        return sb.toString();
    }
}
